import numpy
import scipy.special
import scipy.misc
import npy2cube  # не работает -- надо поменять map на списки
import math

# Подаем на вход квантовые числа
def w(n,l,m,d):
    
    x,y,z = numpy.mgrid[-d:d:30j,-d:d:30j,-d:d:30j]
    
    # Переходим в сферические координаты
    r = lambda x,y,z: numpy.sqrt(x**2+y**2+z**2)  # длина вектора
    theta = lambda x,y,z: numpy.arccos(z/r(x,y,z))  # зенит
    phi = lambda x,y,z: numpy.arctan(y/x)  # азимут
    
    a0 = 1.
    
    # Забыт коэффициент
    k = lambda n,l: numpy.sqrt((2/n/a0)**3 * math.factorial(n-l-1)/(2*n*math.factorial(n+l)))
    
    # Радиальная функция
    R = lambda r,n,l: (2*r/n/a0)**l * numpy.exp(-r/n/a0) * scipy.special.genlaguerre(n-l-1,2*l+1)(2*r/n/a0)
    
    # Волновая функция 
    WF = lambda r,theta,phi,n,l,m: k(n, l) * R(r,n,l) * scipy.special.sph_harm(m,l,phi,theta)
    
    # Квадрат модуля волновой функции(absolute почему-то не работает нормально с мнимыми числами)
    absWF = lambda r,theta,phi,n,l,m: numpy.lib.scimath.sqrt(numpy.real(WF(r,theta,phi,n,l,m))**2 - numpy.imag((WF(r,theta,phi,n,l,m))**2))**2

    ### может тут чего то не хватает? # absWF ничего не выводит в pymol, использую WF
    return WF(r(x,y,z),theta(x,y,z),phi(x,y,z),n,l,m)

# определите шаг grid при заданном диапозоне от -d до d
d = 30
step = 1

#Для каждого главного числа n от 1 до 3 перебираем орбитальные числа l от 0 до n-1, и для них перебираем магнитные m от -l до l

for n in range(1,4):
    d = 10 * n
    step = d / 20
    for l in range(0,n):
        for m in range(-l,l+1):
            grid = w(n,l,m,d)
            name='%s_%s_%s' % (n,l,m)
            # для сохранения нужно задать координаты старта grid и шаг по каждому направлению
            npy2cube.npy2cube(grid,(-d,-d,-d),(step,step,step),'cube/'+name+'.cube')

Визуализация

# ramp -- функция переводящая значение WF в цвет в pymol
# строка -- значение, R, G, B, alpha
cmd.volume_ramp_new('ramp007', [-0.015, 0.00, 0.00, 0.00, 0.00, 
      -0.01,  1.00, 0.00, 0.00, 0.20, 
      -0.005, 0.00, 0.00, 1.00, 0.00, 
      0.005, 0.00, 0.00, 1.00, 0.00, 
      0.01,  0.00, 1.00, 0.00, 0.20, 
      0.015, 0.00, 0.00, 0.00, 0.00])

#Для каждой тройки квантовых чисел загружаем соответствующий файл cube и визуализируем его
for n in range(1, 4):
    for l in range(0, n):
        for m in range(-l, l+1):
            name = f'{n}_{l}_{m}'
            cmd.load(f'~/Desktop/Molmod/cube/{name}.cube')
            cmd.volume(f'{name}_v', name, ramp='ramp007')
            cmd.hide()
            cmd.do(f'show volume, {name}_v')
            cmd.reset()
            cmd.turn('x', 45)
            cmd.turn('y', 45)
            cmd.turn('z', 45)

import glob
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

numbers = ['1_0_0', '2_0_0', '2_1_0', '2_1_1', '2_1_-1', '3_0_0', '3_1_0', '3_1_1', '3_1_-1', '3_2_0', '3_2_1', '3_2_-1', '3_2_2', '3_2_-2']
img_paths = [f'cube/{i}.png' for i in numbers]
images = []
for img_path in img_paths:
    images.append(mpimg.imread(img_path))

orbitals = ['s1','s2','p2','p2','p2','s3','p3','p3','p3','d3','d3','d3','d3','d3']
f = plt.figure(figsize=(20,10))
columns = 4
for i, image in enumerate(images):
    f.add_subplot(len(images) / columns + 1, columns, i + 1)
    plt.imshow(image)
    plt.xticks([])
    plt.yticks([])
    plt.title(f'{orbitals[i]}')

plt.show()

# файл для orca
'''
! UHF SVP XYZFile
%plots 
Format Cube
 MO("H-1.cube",1,0);
 MO("H-2.cube",2,0);
 MO("H-3.cube",3,0);
 MO("H-4.cube",4,0);
 MO("H-5.cube",5,0);
 MO("H-6.cube",6,0);
 MO("H-7.cube",7,0);
 MO("H-8.cube",8,0);
end
* xyz 0 4
 H 0 0 0
*
'''

# ramp -- функция переводящая значение WF в цвет в pymol
# строка -- значение, R, G, B, alpha
cmd.volume_ramp_new('ramp007', [-0.015, 0.00, 0.00, 0.00, 0.00, 
      -0.01,  1.00, 0.00, 0.00, 0.20, 
      -0.005, 0.00, 0.00, 1.00, 0.00, 
      0.005, 0.00, 0.00, 1.00, 0.00, 
      0.01,  0.00, 1.00, 0.00, 0.20, 
      0.015, 0.00, 0.00, 0.00, 0.00])

#Для каждой выдачи orca загружаем соответствующий файл cube, визуализируем его и сохраняем
for n in range(1, 9):
            name = f'H-{n}'
            cmd.load(f'~/Desktop/Molmod/cube/{name}.cube')
            cmd.volume(f'{name}_v', name, ramp='ramp007')
            cmd.hide()
            cmd.do(f'show volume, {name}_v')
            cmd.reset()
            cmd.turn('x', 45)
            cmd.turn('y', 45)
            cmd.turn('z', 45)

import glob
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

numbers = ['H-1', 'H-2', 'H-3', 'H-4']
img_paths = [f'cube/{i}.png' for i in numbers]
images = []
for img_path in img_paths:
    images.append(mpimg.imread(img_path))

orbitals = ['s1','s2','p2','p2','p2','s3','p3','p3','p3','d3','d3','d3','d3','d3']
f = plt.figure(figsize=(20,10))
columns = 4
for i, image in enumerate(images):
    f.add_subplot(len(images) / columns + 1, columns, i + 1)
    plt.imshow(image)
    plt.xticks([])
    plt.yticks([])
    plt.title(f'{orbitals[i]}')

plt.show()

Орбитали полученные с помощью orca похожи на полученные с помощью скрипта, но направлены по-другому (хотя относително друг друга направлены одинаково), к тому же в первом случае на p орбитали неправильно показаны магнитные числа. Еще почему-то скрипт для orca не рисует больше орбиталей после. H-4